import snapshot
import scipy
import scipy.interpolate
import particles as p
import numpyUtil as nu
from numpy import *

import pyx
from pyx.graph import *

def totalMassEnclosed(dat, **kw):
    return massFunc(concatenate((dat.rs, dat.rd, dat.rg)),
                    concatenate((dat.ms, dat.md, dat.mg)), **kw)

def totalProjectedMassEnclosed(dat, **kw):
    return massFunc(concatenate((dat.rs[:,0:2], dat.rd[:,0:2], dat.rg[:,0:2])),
                    concatenate((dat.ms, dat.md, dat.mg)), **kw)

def massFunc(*a, **kw):
    ms, rbins = [concatenate(([0],arr))
                 for arr in p.massProfile(*a, **kw)]
    f = scipy.interpolate.interp1d(rbins, ms)
    def g(x):
        if isscalar(x):
            return f(x)[0]
        return f(x)                    
    return g

def threeDPointMass(dat, **kw):
    massEnc = totalMassEnclosed(dat, **kw)
    
    def f(r):        
        return sqrt(2*snapshot.G*massEnc(r)/r) # FIXME
    return f

def twoDPointMass(dat, **kw):
    massEnc = totalProjectedMassEnclosed(dat, **kw)
    def f(r):        
        return sqrt(2*snapshot.G*massEnc(r)/r)
    return f

def threeDNfw(dat):
    massEnc = totalMassEnclosed(dat)

    # FIXME -- take a stab at calculating virial radius, scale radius
    M = dat.ms.sum() + dat.md.sum() + dat.mg.sum()
    sigma = p.mag(dat.vd).mean() # TODO should in principle be mass weighted
    rVir = snapshot.G*M/sigma**2
    rs = rVir/10
    def f(r):
        x = r/rs
        return sqrt(2*(snapshot.G*massEnc(r)/r) * (x + 1 + (x+1)*log(x+1) ) / (1 + log(x+1)))
    return f

def twoDNfw(dat):
    massEnc = totalProjectedMassEnclosed(dat)
    # FIXME -- take a stab at calculating virial radius, scale radius
    M = dat.ms.sum() + dat.md.sum() + dat.mg.sum()
    sigma = p.mag(dat.vd).mean() # TODO should in principle be mass weighted
    rVir = snapshot.G*M/sigma**2
    rs = rVir/10
    def f(r):        
        x = r/rs
        return sqrt(2*(snapshot.G*massEnc(r)/r) * (x + 1 + (x+1)*log(x+1) ) / (1 + log(x+1)))
    return f

def threeDPowerLaw(dat):
    massEnc = totalMassEnclosed(dat)
    def f(r):        
        return sqrt(snapshot.G*massEnc(r)/r)
    return f

def twoDPowerLaw(dat):
    massEnc = totalProjectedMassEnclosed(dat)
    def f(r):        
        return sqrt(snapshot.G*massEnc(r)/r)
    return f

def threeDPointMassBaryons(dat):
    massEnc = massFunc(dat.rs, dat.ms)
    def f(r):        
        return sqrt(2*snapshot.G*massEnc(r)/r)
    return f

def twoDPointMassBaryons(dat):
    massEnc = massFunc(dat.rs[:,0:2], dat.ms)
    def f(r):        
        return sqrt(2*snapshot.G*massEnc(r)/r)
    return f

def threeDNfwBaryons(dat):
    massEnc = massFunc(dat.rs, dat.ms)
    # FIXME -- take a stab at calculating virial radius, scale radius
    M = dat.ms.sum() + dat.md.sum() + dat.mg.sum()
    sigma = p.mag(dat.vd).mean() # TODO should in principle be mass weighted
    rVir = snapshot.G*M/sigma**2
    rs = rVir/10

    def f(r):        
        x=r/rs
        return sqrt(2*(snapshot.G*massEnc(r)/r) * (x + 1 + (x+1)*log(x+1) ) / (1 + log(x+1)))
    return f

def twoDNfwBaryons(dat):
    massEnc = massFunc(dat.rs[:,0:2], dat.ms)
    # FIXME -- take a stab at calculating virial radius, scale radius
    M = dat.ms.sum() + dat.md.sum() + dat.mg.sum()
    sigma = p.mag(dat.vd).mean() # TODO should in principle be mass weighted
    rVir = snapshot.G*M/sigma**2
    rs = rVir/10

    def f(r):        
        x=r/rs
        return sqrt(2*(snapshot.G*massEnc(r)/r) * (x + 1 + (x+1)*log(x+1) ) / (1 + log(x+1)))
    return f

def threeDPowerLawBaryons(dat):
    massEnc = massFunc(dat.rs, dat.ms)
    def f(r):        
        return sqrt(snapshot.G*massEnc(r)/r)
    return f

def twoDPowerLawBaryons(dat):
    massEnc = massFunc(dat.rs[:,0:2], dat.ms)
    def f(r):        
        return sqrt(snapshot.G*massEnc(r)/r)
    return f

def vescPlot(dat, props, vesc=None, filename=None):
    if vesc is None: vesc = ThreeDPointMass(dat)
    dat.toBodyCoordinates(props)

    # Choose a particle
    # Choose angles s.t. cos(th) has flat dist
    # Map to projected radii
    # plot vesc_real vs. vesc_calc for 
    nPoints = 50
    nAngles = 100
    idxs = array(len(dat.rs)*random.rand(nPoints), 'i')
    rs = p.mag(dat.rs[idxs])
    ths = arccos(random.rand(nAngles))
    # rps[point, angle]
    rps = rs[:,newaxis]*sin(ths[newaxis,:])
    vescActual = sqrt(-2*dat.sPotential[idxs])
    shape = rps.shape
    vescApparent = vesc(rps.ravel()).reshape(shape)
    # now plot them
    points = []
    for vescApp, vescAct in zip(vescApparent, vescActual):
        for ve in vescApp:
            points.append((vescAct, ve))
    
    g=graphxy(width=10,height=10,
              x=axis.log(title='Actual Vesc',min=200,max=2000),
              y=axis.log(title='Predicted Vesc',min=200,max=2000))
    g.plot(data.points(points,x=1,y=2), [style.symbol(symbol=style.symbol.circle,
                                                    size=.01)])
    g.plot(data.function('y(x)=x'))
    if filename: g.writeEPSfile(filename)
    return g

def rpVsVescPlot(dat, props, vesc=None, filename=None):
    dat.toBodyCoordinates(props)

    # Choose a particle
    # Choose angles s.t. cos(th) has flat dist
    # Map to projected radii
    # plot vesc_real vs. vesc_calc for 
    nPoints = 50
    nAngles = 50
    idxs = array(len(dat.rs)*random.rand(nPoints), 'i')
    rs = p.mag(dat.rs[idxs])
    ths = arccos(random.rand(nAngles))
    rpsList = rs[:,newaxis]*sin(ths[newaxis,:])
    vescActual = sqrt(-2*dat.sPotential[idxs])
    # now plot them
    points = []
    for rps, vescAct in zip(rpsList, vescActual):
        for rp in rps:
            points.append((rp, vescAct))
    
    g=graphxy(width=10,
              x=axis.log(title='Projected Radius',min=0.1,max=50),
              y=axis.log(title='Actual Vesc',min=100, max=2000))
    g.plot(data.points(points,x=1,y=2), [style.symbol(symbol=style.symbol.circle,
                                                    size=.01)])
    vescLines(dat, g)
    if filename: g.writeEPSfile(filename)
    return g

def rVsVescPlot(dat, props, vesc=None, filename=None):
    dat.toBodyCoordinates(props)

    # Choose a particle
    # Choose angles s.t. cos(th) has flat dist
    # Map to projected radii
    # plot vesc_real vs. vesc_calc for 
    nPoints = 1000
    idxs = array(len(dat.rs)*random.rand(nPoints), 'i')
    rs = p.mag(dat.rs[idxs])
    vescActual = sqrt(-2*dat.sPotential[idxs])
        
    g=graphxy(width=10,
              x=axis.log(title='r (kpc)',min=0.1,max=50),
              y=axis.log(title='Actual Vesc',min=100, max=2000),
              key=key.key(pos='bl'))
    g.plot(data.array([rs, vescActual],x=1,y=2,title=None),
           [style.symbol(symbol=style.symbol.circle, size=.01)])
    vescLines(dat, g)
    if filename: g.writeEPSfile(filename)
    return g

def vescLines(s, g):
    xs = lnrange(0.1, 100, 50)
    pdata1 = [ data.array([xs, vesc(xs)], x=1,y=2,title=title)
              for vesc,title in ((threeDNfw(s), 'NFW'),
                                 (threeDPointMass(s), 'Point Mass'),
                                 (threeDPowerLaw(s), 'Power Law'))]
    pdata2 = [ data.array([xs, vesc(xs)], x=1,y=2,title=None)
              for vesc in (threeDNfwBaryons(s),
                                 threeDPointMassBaryons(s),
                                 threeDPowerLawBaryons(s))]
    g.plot(pdata1, [style.line()])
    g.plot(pdata2, [style.line(lineattrs=[pyx.color.rgb.red])])
    return g
              

## r = G m / s^2 

### What is a good approx for escape velocity?
